/**
 * TensorFlow模型加载和缓存相关工具函数
 */
import * as tf from "@tensorflow/tfjs"
import * as cocoSsd from "@tensorflow-models/coco-ssd"
// 自定义模型URL（可设置为自己服务器上的模型地址）
export const CUSTOM_MODEL_URL = "/models/coco-ssd/model.json"
// 敏感对象类型定义（COCO-SSD中包含的我们需要检测的敏感对象类型，主要针对身份证上的人脸）
export const sensitiveObjects = ["person", "face", "id_card", "tie"]
/**
 * 从IndexedDB缓存加载模型
 * @param onProgress 加载进度回调函数
 * @returns Promise，解析为模型对象或null
 */
export const loadModelFromCache = async (onProgress?: (message: string) => void): Promise<cocoSsd.ObjectDetection | null> => {
    try {
        // 检查模型是否已缓存
        const models = await tf.io.listModels()
        if (models["indexeddb://coco-ssd-model"]) {
            if (onProgress) onProgress("正在从缓存加载模型...")
            return await cocoSsd.load({
                modelUrl: "indexeddb://coco-ssd-model",
            })
        }
        return null
    } catch (error) {
        console.log("缓存模型获取失败，将从CDN加载", error)
        return null
    }
}
/**
 * 缓存模型到IndexedDB
 * @param loadedModel 已加载的模型对象
 * @returns Promise，解析为void
 */
export const saveModelToCache = async (loadedModel: cocoSsd.ObjectDetection): Promise<void> => {
    try {
        // @ts-ignore - TensorFlow.js类型定义可能不完整
        await loadedModel.model.save("indexeddb://coco-ssd-model")
        console.log("模型已缓存到IndexedDB")
    } catch (error) {
        console.error("模型缓存失败:", error)
    }
}
/**
 * 加载COCO-SSD模型
 * @param onProgress 进度回调函数，接收消息和进度百分比
 * @returns Promise，解析为加载好的模型
 */
export const loadModel = async (onProgress: (message: string, progress?: number) => void): Promise<cocoSsd.ObjectDetection> => {
    try {
        // 设置TensorFlow.js后端为WebGL以提高性能
        await tf.setBackend("webgl")
        onProgress("正在加载模型，首次加载可能需要几分钟...")

        // 尝试从缓存加载
        let cachedModel = await loadModelFromCache(onProgress)
        if (cachedModel) {
            console.log("从缓存加载模型成功")
            return cachedModel
        }

        // 监听加载进度
        let lastProgress = 0

        // 监听进度
        try {
            // @ts-ignore - TensorFlow.js类型定义可能不完整
            tf.env().set("TENSORFLOWJS_PROGRESS_PLUGIN_URL_FILTER", /\.bin/)

            const progressHandler = (fraction: number) => {
                // 避免进度条回退
                if (fraction > lastProgress) {
                    lastProgress = fraction
                    const progressPercent = Math.floor(fraction * 100)
                    onProgress(`正在加载模型... ${progressPercent}%`, progressPercent)
                }
            }

            // 注册进度回调
            // @ts-ignore - TensorFlow.js类型定义可能不完整
            tf.env().set("DOWNLOAD_PROGRESS_HANDLER", progressHandler)
        } catch (err) {
            console.warn("无法设置模型加载进度监听器", err)
        }

        // 从自定义URL加载模型
        const model = await cocoSsd.load({ modelUrl: CUSTOM_MODEL_URL })

        // 缓存模型以供将来使用
        await saveModelToCache(model)

        return model
    } catch (error) {
        console.error("模型加载失败:", error)
        throw new Error("模型加载失败，请刷新页面重试")
    }
}
/**
 * 使用模型检测图像中的敏感区域
 * @param model COCO-SSD模型
 * @param canvas 包含图像的Canvas元素
 * @returns Promise，解析为检测到的敏感区域数组
 */
export const detectSensitiveAreas = async (model: cocoSsd.ObjectDetection, canvas: HTMLCanvasElement): Promise<Array<[number, number, number, number]>> => {
    // 使用模型检测对象
    const predictions = await model.detect(canvas)
    console.log(predictions)
    // 过滤出敏感对象
    const sensitiveRegions = predictions.filter((prediction) => sensitiveObjects.includes(prediction.class)).map((region) => region.bbox)

    return sensitiveRegions
}
